package 算法回顾;

import com.alibaba.fastjson.JSON;

import java.util.*;

/**
 * @description:
 * @author: 小白白
 * @create: 2021-09-20
 **/

public class No692前K个高频单词 {

    public List<String> topKFrequent(String[] words, int k) {

        Map<String,Integer> map = new HashMap<>();
        TreeMap<Integer, HashSet<String>> mapSet = new TreeMap<>((o1, o2)-> o2 - o1);

        for (String word : words) {
            int count = map.getOrDefault(word, 0) + 1;
            map.put(word, count);

            HashSet<String> set = mapSet.getOrDefault(count - 1, new HashSet<>());
            set.remove(word);
            mapSet.put(count - 1, set);

            HashSet<String> thisSet = mapSet.getOrDefault(count, new HashSet<>());
            thisSet.add(word);
            mapSet.put(count, thisSet);

        }

        List<String> result = new ArrayList<>(k);
        for (Map.Entry<Integer, HashSet<String>> entry : mapSet.entrySet()) {
            HashSet<String> set = entry.getValue();
            if (set.size() <= k) {
                List<String> tempResult = new ArrayList<>(set.size());
                tempResult.addAll(set);
                Collections.sort(tempResult);
                result.addAll(tempResult);
            }
            if (set.size() > k) {
                List<String> tempResult = new ArrayList<>(set.size());
                tempResult.addAll(set);
                Collections.sort(tempResult);
                tempResult = tempResult.subList(0, k);
                result.addAll(tempResult);
            }
            k -= set.size();
            if (k <= 0) {
                break;
            }
        }

        return result;
    }

    public static void main(String[] args) {
        No692前K个高频单词 n = new No692前K个高频单词();
        String[] words = {"the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"};
        List<String> result = n.topKFrequent(words, 4);
        System.out.println(JSON.toJSONString(result));
    }

}
